#pragma once
#include "AlgoEddyTrack.h"
#include <vector>
#include <string>
#include <memory>
#include <utility>

namespace EddyDetection{
    // 2D Kalman Filter
    class KalmanFilter2D {
    public:
        KalmanFilter2D(const BoundaryPoint2D& initial_pos, double dt = 1.0);
        void predict();
        void update(const BoundaryPoint2D& measurement);
        BoundaryPoint2D get_position() const;
        BoundaryPoint2D get_velocity() const;
        BoundaryPoint2D predict_next_position() const;
        void set_process_noise(double process_noise);
        void set_measurement_noise(double measurement_noise);

        // State setting methods
        void set_state(const BoundaryPoint2D& position, const BoundaryPoint2D& velocity);
        void set_position(const BoundaryPoint2D& position);
        void set_velocity(const BoundaryPoint2D& velocity);

        // Get covariance matrix for uncertainty assessment
        const std::vector<std::vector<double>>& get_covariance_matrix() const;

    private:
        // State vector [lat, lon, vlat, vlon]
        std::vector<double> x;
        // State covariance matrix
        std::vector<std::vector<double>> P;
        // State transition matrix
        std::vector<std::vector<double>> F;
        // Observation matrix
        std::vector<std::vector<double>> H;
        // Process noise covariance matrix
        std::vector<std::vector<double>> Q;
        // Measurement noise covariance matrix
        std::vector<std::vector<double>> R;
        // Time step
        double dt;

        // Matrix operation helper functions
        std::vector<std::vector<double>> matrix_multiply(
            const std::vector<std::vector<double>>& A,
            const std::vector<std::vector<double>>& B) const;
        std::vector<std::vector<double>> matrix_add(
            const std::vector<std::vector<double>>& A,
            const std::vector<std::vector<double>>& B) const;
        std::vector<std::vector<double>> matrix_subtract(
            const std::vector<std::vector<double>>& A,
            const std::vector<std::vector<double>>& B) const;
        std::vector<std::vector<double>> matrix_transpose(
            const std::vector<std::vector<double>>& A) const;
        std::vector<std::vector<double>> matrix_inverse(
            const std::vector<std::vector<double>>& A) const;
        void initialize_matrices();
    };


    // 2D Eddy Tracker
    class EddyTracker {
    public:
        EddyTracker(double max_match_dist = 50.0, double dt = 1.0);  // Maximum matching distance, time step size
        void init_tracks(const Frame& first_frame);
        void match_and_update(const Frame& frame);
        void kalman_update(const Frame& frame);
        void handle_split_merge();
        std::vector<Track> get_tracks() const;
        void set_kalman_parameters(double process_noise, double measurement_noise);

        // Split-merge related parameter settings
        void set_split_merge_parameters(double merge_distance_ratio = 0.3,
            double split_area_threshold = 2.0,
            double merge_area_threshold = 0.5);

    private:
        std::vector<Track> tracks;
        // Maintain independent 2D Kalman filters for each trajectory
        std::vector<std::unique_ptr<KalmanFilter2D>> kalman_filters;
        int next_track_id = 0;
        double max_match_dist;
        double dt; // Time step
        double process_noise = 1.0;
        double measurement_noise = 10.0;

        // Split-merge parameters
        double merge_distance_ratio = 0.3;    // Merge distance threshold ratio
        double split_area_threshold = 2.0;    // Split area threshold
        double merge_area_threshold = 0.5;    // Merge area threshold

        // Helper functions
        void add_new_track_with_kalman(const Frame& frame, const EddyContourData& eddy, int layer_idx);
        void update_track_with_kalman(size_t track_idx, const EddyContourData& eddy, int frame_idx);
        void remove_dead_tracks(const Frame& frame);

        // Split-merge detection helper functions
        void detect_and_handle_merges(const Frame& current_frame);
        void detect_and_handle_splits(const Frame& current_frame);
        bool check_merge_criteria(size_t track_i, size_t track_j, const Frame& current_frame);
        bool check_split_criteria(size_t track_idx, const Frame& current_frame,
            std::vector<std::pair<int, int>>& potential_splits);
        void execute_merge(size_t keep_track, size_t remove_track, double confidence_keep, double confidence_remove);
        void execute_split(size_t parent_track, const std::vector<std::pair<int, int>>& split_eddy_indices,
            const Frame& current_frame);

        // Trajectory quality assessment
        double calculate_track_confidence(size_t track_idx) const;
        double calculate_prediction_error(size_t track_idx, const BoundaryPoint2D& actual_pos) const;

        // Geographic distance calculation using Haversine formula
        double haversine_distance(const BoundaryPoint2D& p1, const BoundaryPoint2D& p2) const;

    public:
        // Public method to access haversine distance calculation
        static double calculate_haversine_distance(const BoundaryPoint2D& p1, const BoundaryPoint2D& p2);
    };

    // Calculate 2D Euclidean distance (for non-geographic coordinates)
    double euclidean_distance_2d(const BoundaryPoint2D& p1, const BoundaryPoint2D& p2);

}
